from openpyxl import Workbook, load_workbook

# 定义ExcelOperator类
class ExcelOperator:
    def __init__(self, file_path):
        """
        初始化方法，用于加载Excel文件。

        参数：
        file_path (str): Excel文件的路径，如果文件不存在则会创建一个新的工作簿。
        """
        self.file_path = file_path
        try:
            self.workbook = load_workbook(self.file_path)
            self.sheet = self.workbook.active
        except FileNotFoundError:
            self.workbook = Workbook()
            self.sheet = self.workbook.active

    # 定义获取单元格数据
    def get_cell_value(self, cell_address):
        """
        获取指定单元格的数据。

        参数：
        cell_address (str): 单元格地址，例如 "A1"。

        返回：
        返回指定单元格中的数据，如果单元格为空则返回None。
        """
        return self.sheet[cell_address].value
    
    # 定义对单元格写入数据的方法
    def set_cell_value(self, cell_address, value):
        """
        对指定单元格写入数据。

        参数：
        cell_address (str): 单元格地址，例如 "A1"。
        value (any): 要写入单元格的数据，可以是各种Python基本数据类型（如字符串、数字、布尔值等）。
        """
        self.sheet[cell_address] = value
        self.workbook.save(self.file_path)

    # 定义选择工作表的方法
    def select_sheet(self, sheet_name):
        """
        选择工作表。

        参数：
        sheet_name (str): 工作表名称。
        """
        self.sheet = self.workbook[sheet_name]

    # 新增方法，用于将数据批量写入指定区域
    def write_data_to_area(self, start_cell_address, data):
        """
        将给定的数据批量写入到Excel工作表的指定区域。

        参数：
        start_cell_address (str): 起始单元格地址，例如 "A1"，代表要写入数据区域的左上角单元格。
        data (list of lists): 要写入的数据，是一个二维列表，内层列表代表每一行的数据。

        示例：
        若 data = [[1, 2], [3, 4]]，start_cell_address = "A1"，则会将数据按行依次写入从A1开始的区域。
        """
        row_index = 0
        col_index = 0
        start_row, start_col = self._parse_cell_address(start_cell_address)

        for row_data in data:
            for cell_data in row_data:
                cell_address = self._generate_cell_address(start_row + row_index, start_col + col_index)
                self.set_cell_value(cell_address, cell_data)
                col_index += 1
            row_index += 1
            col_index = 0
   
    # 新增方法，用于获取指定区域的数据
    def get_data_from_area(self, start_cell_address, end_cell_address):
        """
        获取Excel工作表中指定区域的数据，并以二维列表形式返回。

        参数：
        start_cell_address (str): 起始单元格地址，例如 "A1"，代表要获取数据区域的左上角单元格。
        end_cell_address (str): 结束单元格地址，例如 "C3"，代表要获取数据区域的右下角单元格。

        返回：
        返回一个二维列表，内层列表代表每一行获取到的数据。
        """
        start_row, start_col = self._parse_cell_address(start_cell_address)
        end_row, end_col = self._parse_cell_address(end_cell_address)

        result_data = []
        for row in range(start_row, end_row + 1):
            row_data = []
            for col in range(start_col, end_col + 1):
                cell_address = self._generate_cell_address(row, col)
                cell_value = self.get_cell_value(cell_address)
                row_data.append(cell_value)
            result_data.append(row_data)

        return result_data

    # 用于解析单元格地址字符串
    def _parse_cell_address(self, cell_address):
        """
        将单元格地址字符串解析为行号和列号（数字形式）。

        参数：
        cell_address (str): 单元格地址，例如 "A1"。

        返回：
        返回解析后的行号和列号（均为整数），例如 (1, 1) 对应 "A1"。
        """
        col = 0
        col_str = ""
        for char in cell_address:
            if char.isalpha():
                col_str += char
            else:
                break
        for i, char in enumerate(col_str[::-1]):
            col += (ord(char.upper()) - ord('A') + 1) * (26 ** i)
        row = int(cell_address[len(col_str):])
        return row, col

    # 用于生成单元格地址字符串
    def _generate_cell_address(self, row, col):
        """
        根据给定的行号和列号生成单元格地址字符串。

        参数：
        row (int): 行号，从1开始计数。
        col (int): 列号，从1开始计数。

        返回：
        返回生成的单元格地址字符串，例如 "A1"。
        """
        col_str = ""
        while col > 0:
            col, remainder = divmod(col - 1, 26)
            col_str = chr(ord('A') + remainder) + col_str
        return col_str + str(row)


# 封装测试函数
def test_write_data_to_area():
    """
    测试write_data_to_area方法是否能正确将数据写入指定区域的测试函数。
    """
    file_path = "test.xlsx"
    sheet_name = "Sheet1"
    start_cell_address = "A1"
    data_to_write = [
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]
    ]

    # 创建ExcelOperator实例
    excel_op = ExcelOperator(file_path)

    # 选择工作表
    excel_op.select_sheet(sheet_name)

    # 调用write_data_to_area方法写入数据
    excel_op.write_data_to_area(start_cell_address, data_to_write)

    # 验证写入的数据是否正确
    for row_index, row_data in enumerate(data_to_write):
        for col_index, expected_value in enumerate(row_data):
            cell_address = excel_op._generate_cell_address(row_index + 1, col_index + 1)
            actual_value = excel_op.get_cell_value(cell_address)
            assert actual_value == expected_value, f"数据验证失败，预期值为 {expected_value}，实际值为 {actual_value}，单元格地址为 {cell_address}"

    print("测试通过，write_data_to_area方法功能正常。")


if __name__ == "__main__":
    # test_write_data_to_area()
    pass